KV Cache 与推理优化

标签

AI/llm

字数

4299 字

阅读时间

18 分钟

自回归解码 每次只生成一个 token，每个新 token 都要对之前所有 token 做注意力。天真的做法是每一步把整段前缀重新前向一遍——计算量按序列长度平方增长。

这条路走不通。而绕开它的那个东西（KV Cache）决定了三件事：推理多快、显存占多少、能支持多长上下文和多高并发。vLLM 的 PagedAttention、SGLang 的 RadixAttention，本质都是在跟它抢内存。

为什么能缓存：谁变了，谁没变

看注意力在解码时的形态：

• 新 token 的 Q 是新的——它在问「我应该关注谁」
• 历史 token 的 K 和 V 是固定的——它们不随新 token 变化

所以最朴素也最有效的优化就是：把历史 token 的 K 和 V 存起来，下次直接用。

每一步的注意力复杂度从 $O(n^2)$ 降到 $O(n)$。 代价是缓存要占一块显存，而且随序列长度线性增长。

KV Cache 的本质是用显存换算力。 这也直接解释了「长上下文 = 大显存」。

两个阶段的性格完全不同

推理被拆成 Prefill 和 Decode，两者的资源特征几乎相反：

	Prefill	Decode
每次前向的输入	整个 prompt（P 个 token）	1 个 token
缓存动作	填充全部 P 个位置、全部层	追加 1 个位置
矩阵形状	matrix–matrix（Q 有 P 行）	matrix–vector（单行 Q）
瓶颈	计算受限（GPU FLOPs）	显存带宽受限
并行性	所有 prompt 位置并行	本质串行
决定什么延迟	TTFT（首 token 时间）	TPOT（每输出 token 时间）

Decode 为什么是带宽受限：它是细的矩阵-向量乘，算术强度低——要把整个模型权重（外加不断增长的 KV cache）从 HBM 流一遍，却只做很少的计算。

由此推出服务系统的第一个设计：必须把多个并发 decode 批在一起，把「一次性读权重」的开销摊到多个请求上。这也是为什么 batch size 在 decode 阶段比在 prefill 阶段重要得多。

第二个设计：大部署会把两个阶段分开调度甚至放到不同硬件上（disaggregated prefill/decode）——因为它们的瓶颈在两边，混在一起调度会互相等待。

这条和 Context Engineering 里那套 cold prefill / resume prefill / short decode 的分解是同一件事：那里是从成本角度拆，这里是从硬件瓶颈角度拆，指的是同两个阶段。

两个阶段的矩阵形状不同，瓶颈也就落在了不同的资源上。

Prefill（一次处理整个 prompt，P 个 token）
    Q (P × d)  @  K^T (d × P)  ──▶  P × P 的分数矩阵
    ├─ matrix–matrix：Q 有 P 行，GPU 一次吃满
    ├─ 瓶颈：计算（FLOPs）
    └─ 决定 TTFT

Decode（每步 1 个 token）
    Q (1 × d)  @  K^T (d × n)  ──▶  1 × n 的分数向量
    ├─ matrix–vector：要把整个模型权重（外加不断增长的 KV）从 HBM 流一遍，
    │  却只做很少的计算 ──▶ 算术强度低，卡在带宽上
    ├─ 本质串行，靠并行填不满算力
    └─ 决定 TPOT

由此推出服务系统的两个设计
① 把多个并发 decode 批在一起 ──▶ 一次读权重的开销摊给多个请求
   （batch size 在 decode 阶段比在 prefill 阶段重要得多，原因在这里）
② 大部署把两个阶段分开调度，甚至放到不同硬件上 ——
   它们的瓶颈在两边，混在一起调度只会互相等待

KV Cache 有多大：一张显存账单

单 token 的缓存字节数：

$$\text{bytes\_per\_token}=2\times \text{layers}\times \text{kv\_heads}\times \text{head\_dim}\times \text{bytes\_per\_element}$$

因子	含义
2	K 和 V 各一份
layers	Transformer 层数
kv_heads	KV 头数——注意不是 query heads
head_dim	每个头的维度
bytes_per_element	FP16/BF16 是 2 字节；INT8/FP8 是 1 字节

kv_heads 是承重变量。 这就是为什么**「模型大小」单独不足以预测 KV 压力**——两个总参数量相近的模型，KV 占用可能差好几倍，只看参数量看不出来。

各模型的实测对照（FP16，每 token）

模型	注意力布局	KV heads	每 token KV
Llama-1 65B	MHA	64	2560 KiB
Llama-2 / Llama-3 70B	GQA	8	320 KiB
Llama-3 8B	GQA	8	128 KiB

Llama-1 65B 和 Llama-3 70B 参数量接近，但每 token 的 KV 差了 8 倍——差别全在布局上。

一个具体算例

7B 量级（约 40 层、n_kv_heads=32、head_dim=128），FP16，单请求：

序列长度	KV Cache
2,048	1.34 GB
32K	约 21 GB
并发 8 请求	再 ×8

32K 上下文的 21 GB 几乎吃掉一张 A100 的一半。

实测的经验值：KV Cache 能占到总显存的 50%–80%，模型权重反而只占小头。

所以「显存不够」很多时候卡在 KV Cache 上，瓶颈不在模型权重。

三种注意力布局：KV Cache 的轻量化路径

KV 缓存的大小由 KV 头数决定，而不是 query 头数——布局选择就在利用这个区别：

布局	KV heads	权衡
MHA（多头注意力）	= query heads（如 64）	表达能力最强，缓存最大
GQA（分组查询注意力）	中间值（64 query / 8 KV，8:1 分组）	缓存缩到 1/8，质量影响很小 → 现代主流默认
MQA（多查询注意力）	1	缓存最小，但过度压缩导致表达能力大幅衰减、生成质量明显下降

演进的方向很清楚：MHA → GQA → 更极致的压缩（如 MLA），每一代都在解决落地时的同一个痛点——显存与带宽。

GQA 成为默认这件事本身就是一条判据：它说明在 KV 缓存这个维度上，8 倍压缩的收益远大于那点质量损失。MQA 压缩得更狠却没能普及，说明这条曲线在 1 这个位置已经过头了。

KV Cache 的大小由 KV 头数决定，布局选择就在拿 KV 头数换缓存。

MHA：KV 头数 = Q 头数
   Q1  Q2  Q3  Q4  Q5  Q6  Q7  Q8
    │   │   │   │   │   │   │   │
   K1  K2  K3  K4  K5  K6  K7  K8
        └─ 每组 K/V 独立 ──▶ 缓存最大，表达能力最强

GQA：每 h/g 个 Q 头共享一组 KV（下图按 4 个 Q 头共享一组示意）
   Q1  Q2  Q3  Q4   Q5  Q6  Q7  Q8
   └───┴───┴───┘    └───┴───┴───┘
        K1               K2
        └─ 缓存缩到 1/g。Llama-2/3 70B：64 → 8，每 token 的 KV 从 2560 KiB 降到 320 KiB

MQA：KV 头数 = 1
   Q1  Q2  Q3  Q4  Q5  Q6  Q7  Q8
   └───┴───┴───┴───┴───┴───┴───┘
                 K1
        └─ 缓存最小，但表达力衰减明显 ──▶ 压缩过头，没成主流

g = h 时 GQA 退化为 MHA，g = 1 时退化为 MQA —— GQA 是两者的连续折中。
MLA 走的是另一条路：不减少头数，把 K/V 压到低维潜在空间（DeepSeek-V2 报约 93.3% 压缩率）。

服务侧的两个核心优化

PagedAttention：把 KV Cache 当虚拟内存管

朴素做法的三个致命伤（按「最大可能长度」一次性申请一整块连续显存）：

伤	表现
内部碎片	实际只生成了 500 token，剩下的 7000+ 个位置白占着
无法共享	两个请求有相同的系统提示词，却各存了一份
无法扩容	长度预估错了就得整块搬移

vLLM 的思路是：这像极了操作系统的内存碎片问题——那就按操作系统的方式分页。

做法：把 KV Cache 切成固定大小的块（vLLM 默认 B=16 个 token 一块）；每个请求维护一张块表（block table），记录逻辑块 → 物理块的映射；逻辑上连续，物理上可以不连续。

逻辑视角（请求看到的是连续序列，vLLM 默认 B = 16 个 token 一块）
   [ 0–15 ] [ 16–31 ] [ 32–47 ] [ 48–… ]
       │         │         │
       ▼         ▼         ▼
块表（每个请求一张，记录逻辑块 → 物理块的映射）
       块 0 ──▶ 物理块 7
       块 1 ──▶ 物理块 3
       块 2 ──▶ 物理块 12
       │         │         │
       ▼         ▼         ▼
物理块池（显存里等大的块，物理上可以不连续）
   ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐
   │ 0  │ │ 1  │ │ 2  │ │ 3  │ │ …  │ │ 12 │ …
   └────┘ └────┘ └────┘ └────┘ └────┘ └────┘

注意力核按块表去对应物理块收集 K/V，不再要求内存连续。

收益	说明
内部碎片归零	按需分配，容量上限从「最大长度」变成「池总量」
共享前缀	相同前缀的块只存一份，多个请求引用同一物理块（分歧时写时复制）
吞吐提升	该文数据 2–4 倍，显存利用率接近最优

APC：连计算都省了

Automatic Prefix Caching：给每个 block 算一个哈希（块内容的指纹）。新请求进来时，如果它的前缀与之前某个请求完全一致（比如同一段系统提示词），直接复用那个物理块——连 Prefill 都省了，不用重新计算这些 token 的 KV。

for i, block_tokens in enumerate(chunk(request_tokens, 16)):
    if hash(block_tokens) in prefix_cache:  hit_blocks += 1
    else:  break     # 首次失配，后面不再可能命中

两者的分工要分清：

PagedAttention 省的是「内存」，APC 省的是「计算 + 内存」。

vLLM 0.6+ 默认开启 APC。

RadixAttention：把前缀复用做到任意长度

SGLang 的答案是把所有请求的历史 KV 组织成一棵前缀树（Radix Tree），从而做到任意长度的前缀复用——而 APC 只能按完整块复用。

为什么值得：真实业务里大量 KV 是重复的——系统提示词（所有请求都带）、few-shot 示例（同一 agent 任务几乎一样）、对话历史（第 N 轮的请求包含前 N-1 轮全部内容）、结构化生成（JSON schema、工具定义，每个请求重复）。

这条对 Context Engineering 是个直接的呼应：那里说「把静态内容放在提示开头」是为了命中 Prompt Cache；这里说同一件事在服务侧还有第二个收益——前缀树复用连 Prefill 的计算都能省掉。「静态内容放开头」这一条同时吃到两层优化。

三个机制省的东西不一样，别混着记。

PagedAttention（vLLM）
   把 KV Cache 切成固定大小的块 + 一张块表，逻辑连续、物理可以不连续
   ├─ 内部碎片归零：容量上限从「最大长度」变成「池总量」
   ├─ 相同前缀的块只存一份，出现分歧时写时复制
   └─ 省的是「内存」        该文报的吞吐提升 2–4 倍

APC（Automatic Prefix Caching，vLLM 0.6+ 默认开）
   给每个 block 算一个内容哈希，新请求的前缀块命中就直接复用物理块
   ├─ 这些 token 的 Prefill 连算都不用算
   └─ 省的是「计算 + 内存」
        限制：只能按「完整块」复用

RadixAttention（SGLang）
   把所有请求的历史 KV 组织成一棵前缀树（Radix Tree）
   └─ 任意长度的前缀复用，不受块边界限制，省的是「计算 + 内存」

真实业务里重复的 KV 来自四处：系统提示词（所有请求都带）、
few-shot 示例（同一 agent 任务几乎一样）、对话历史（第 N 轮包含前 N−1 轮全部内容）、
结构化生成（JSON schema 与工具定义每个请求重复）。

五个必须知道的坑

坑	说明
一、显存增长主导	长上下文下缓存能远大于权重。生产的吞吐通常受 KV 显存限制，而不是 FLOP 限制
二、变长 / ragged batching	一个 batch 里序列长度不同、结束时间不同。天真地按 max_seq_len 预留会浪费大部分缓存；完成的序列会留下空洞，需要簿记回收 —— 这正是上面 PagedAttention 的动机
三、位置编码一致性	RoPE 的旋转是在缓存 K 之前按绝对位置施加的，所以缓存的 key 已经带上了位置相位。按正确索引旋转就没问题；坑出现在位置发生移动时（滑窗或缓存驱逐）——天真的驱逐会让位置数学错位，除非重新锚定或改用相对方案（见 09-位置编码）
四、打包 batch 的掩码	因果掩码必须让每个 token 只关注自己序列中更早的位置。多序列打包进一个张量时，需要按序列的块对角掩码，否则 token 会跨序列泄漏
五、束搜索 / 多采样	会把缓存乘以 beam width / 采样数，除非共享前缀只存一份、只复制分支

正确性陷阱会让输出直接崩

上面第三、四条一旦出错，后果是乱码或指令遵循失败，比质量下降严重得多。另外还有一条特别隐蔽的：

chat template 的 token 没被包含进缓存前缀。

这条会和 分词那篇里说的「必须用模型自己的 chat template，不要自己拼字符串」叠在一起发作——拼错格式时，缓存前缀也跟着错。

多轮对话里的缓存行为

一条很实用的定义：

前缀不变时，缓存可以跨轮复用；编辑更早的消息，会让那个位置之后的整个后缀失效。

所以聊天类应用里「改上面一条消息」要重算的是它之后的全部，那一条本身只占很小一部分。这也解释了为什么这类编辑操作在长对话里体感上明显更慢。

前缀不变，缓存可以跨轮复用；改动越靠前，失效的后缀越长。

第 1 轮   [system][user1][assistant1]
           └───────┬───────┘
                   └─ 这一段的 KV 被缓存下来

第 2 轮   [system][user1][assistant1][user2]
           └───────┬────────────────────┘
                   └─ 前缀完全一致 ──▶ 直接复用，只需要算 [user2] 这一小段

改动位置决定要重算多少
   改 [system]        ──▶ 全部重算，代价最大
   改 [user1]         ──▶ 从 user1 之后全部重算
      └─ 被改的那一条本身只占很小一部分，代价几乎全落在它后面的后缀上
   改最后一条 [user2]  ──▶ 只影响它自己之后的内容

这就是长对话里「改上面一条消息」体感上明显更慢的原因。

KV 量化

用更低的精度存 K/V（INT8 / FP8 / INT4），但矩阵乘的累加仍保持 BF16/FP16。

收益：缓存占用和读带宽大致减半或减到 1/4，直接提升批容量与 decode 吞吐——因为 decode 就是带宽受限的。

代价是精度，属于「用小质量损失换并发量」的典型取舍。

该监控什么

指标	意义
每个并发请求的 KV 占用量（GB）	判断是显存受限还是算力受限
Prefill 时间 / Decode 时间	分别对应 TTFT 与 TPOT
共享前缀的缓存命中率	直接反映 APC / 前缀树有没有生效

最后一条最容易被忽略——如果系统提示词没被复用，那说明前缀不一致（多一个空格、时间戳变了、chat template 拼法不同），这时候查收益为什么没出现，应该先看命中率。

相关

• 03-Decoder-Only 与自回归 —— KV Cache 要解决的问题从哪来
• 09-位置编码 —— RoPE 与缓存的位置一致性坑
• Context Engineering —— 「静态内容放开头」在成本与服务两侧的双重收益
• 05-文本分词与子词算法：BPE、WordPiece 与 Unigram —— chat template 拼错会让缓存前缀失效
• 01-推理性能指标与瓶颈定位 —— 算术强度与 Roofline 的量化推导、TTFT/TPOT/Goodput 的完整口径
• 02-PagedAttention：KV Cache 的分页管理 —— 本篇讲「这块显存有多大」，那一篇讲「引擎怎么把它管起来」
• 03-推理调度：Continuous Batching 与 Chunked Prefill —— batched decode 的调度实现
• 04-前缀缓存：APC 与 RadixAttention —— APC 块哈希与 RadixAttention 的深入展开
• 08-量化 —— KV Cache 量化（KIVI / FP8）的算法与精度取舍
• 09-Speculative Decoding —— 另一条与 KV Cache 交互的解码优化路线
• 14-长上下文技术 —— KV Cache 随长度线性增长这条约束，在模型侧怎么绕

参考

• http://prachub.com/interview-questions/explain-kv-cache-in-transformer-inference
• https://sampooni.github.io/kv-cache-docs/chapters/ch02-background.html
• https://blog.csdn.net/u013970991/article/details/163229125
• https://aiglossary.xyz/glossary/kv-cache/

贡献者

YJ

文件历史

最后编辑于 5 分钟前查看完整历史

7864e-111

aa2e3-docs：补充

ae1ee-docs：内容完善

38e7a-docs：新增

04687-docs：笔记